package RSA

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"os"
)

func generateKeyPEM() (string, string, error) {
	// 1. 生成RSA密钥对
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return "", "", err
	}
	publicKey := &privateKey.PublicKey

	// 2. 将密钥保存为PEM格式字符串
	privateKeyPEM, err := PrivateKeyToPEMString(privateKey)
	if err != nil {
		return "", "", err
	}

	publicKeyPEM, err := PublicKeyToPEMString(publicKey)
	if err != nil {
		return "", "", err
	}

	return privateKeyPEM, publicKeyPEM, nil
}

// privateKeyToPEMString 将RSA私钥转换为PEM格式字符串
func privateKeyToPEMString(privateKey *rsa.PrivateKey) (string, error) {
	privateKeyBytes := x509.MarshalPKCS1PrivateKey(privateKey)
	privateKeyPEM := pem.EncodeToMemory(
		&pem.Block{
			Type:  "RSA PRIVATE KEY",
			Bytes: privateKeyBytes,
		},
	)
	return string(privateKeyPEM), nil
}

// publicKeyToPEMString 将RSA公钥转换为PEM格式字符串
func publicKeyToPEMString(publicKey *rsa.PublicKey) (string, error) {
	publicKeyBytes, err := x509.MarshalPKIXPublicKey(publicKey)
	if err != nil {
		return "", err
	}
	publicKeyPEM := pem.EncodeToMemory(
		&pem.Block{
			Type:  "PUBLIC KEY",
			Bytes: publicKeyBytes,
		},
	)
	return string(publicKeyPEM), nil
}

// privateKeyFromPEMString 从PEM字符串解析RSA私钥
func privateKeyFromPEMString(privateKeyPEM string) (*rsa.PrivateKey, error) {
	block, _ := pem.Decode([]byte(privateKeyPEM))
	if block == nil {
		return nil, errors.New("failed to parse PEM block containing the private key")
	}

	privateKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return nil, err
	}

	return privateKey, nil
}

// publicKeyFromPEMString 从PEM字符串解析RSA公钥
func publicKeyFromPEMString(publicKeyPEM string) (*rsa.PublicKey, error) {
	block, _ := pem.Decode([]byte(publicKeyPEM))
	if block == nil {
		return nil, errors.New("failed to parse PEM block containing the public key")
	}

	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, err
	}

	publicKey, ok := pub.(*rsa.PublicKey)
	if !ok {
		return nil, errors.New("not an RSA public key")
	}

	return publicKey, nil
}

func encryptByPublicKey(publicKey *rsa.PublicKey, data []byte) ([]byte, error) {
	encryptedMessage, err := rsa.EncryptPKCS1v15(rand.Reader, publicKey, data)
	if err != nil {
		return nil, fmt.Errorf("加密失败: %v", err)
	}

	return encryptedMessage, nil
}

func decryptByPrivateKey(privateKey *rsa.PrivateKey, data []byte) ([]byte, error) {
	// 使用私钥解密
	decryptedMessage, err := rsa.DecryptPKCS1v15(rand.Reader, privateKey, data)
	if err != nil {
		return nil, fmt.Errorf("解密失败: %v", err)
	}

	return decryptedMessage, nil
}

func generateKeyFile(privateFileName, publicFileName string) error {
	// 生成 RSA 密钥对
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return fmt.Errorf("生成密钥对失败: %v", err)
	}

	return saveKeyFile(privateKey, privateFileName, publicFileName)
}

func saveKeyFile(privateKey *rsa.PrivateKey, privateFileName, publicFileName string) error {
	// 保存私钥到文件
	err := savePrivateKeyToFile(privateKey, privateFileName)
	if err != nil {
		return fmt.Errorf("保存私钥失败: %v", err)
	}

	// 保存公钥到文件
	err = savePublicKeyToFile(&privateKey.PublicKey, publicFileName)
	if err != nil {
		return fmt.Errorf("保存公钥失败: %v", err)
	}

	return nil
}

// 保存私钥到文件
func savePrivateKeyToFile(privateKey *rsa.PrivateKey, filename string) error {
	// 将私钥转换为 DER 格式
	privateKeyBytes := x509.MarshalPKCS1PrivateKey(privateKey)

	// 创建 PEM 块
	privateKeyBlock := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: privateKeyBytes,
	}

	// 创建文件并写入 PEM 块
	file, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer file.Close()

	err = pem.Encode(file, privateKeyBlock)
	if err != nil {
		return err
	}

	return nil
}

// 保存公钥到文件
func savePublicKeyToFile(publicKey *rsa.PublicKey, filename string) error {
	// 将公钥转换为 DER 格式
	publicKeyBytes, err := x509.MarshalPKIXPublicKey(publicKey)
	if err != nil {
		return err
	}

	// 创建 PEM 块
	publicKeyBlock := &pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: publicKeyBytes,
	}

	// 创建文件并写入 PEM 块
	file, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer file.Close()

	err = pem.Encode(file, publicKeyBlock)
	if err != nil {
		return err
	}

	return nil
}

// 从文件加载私钥
func loadPrivateKeyFromFile(filename string) (*rsa.PrivateKey, error) {
	// 读取文件内容
	fileContent, err := os.ReadFile(filename)
	if err != nil {
		return nil, err
	}

	// 解码 PEM 块
	block, _ := pem.Decode(fileContent)
	if block == nil {
		return nil, fmt.Errorf("无法解码 PEM 块")
	}

	// 解析私钥
	privateKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return nil, err
	}

	return privateKey, nil
}

// 从文件加载公钥
func loadPublicKeyFromFile(filename string) (*rsa.PublicKey, error) {
	// 读取文件内容
	fileContent, err := os.ReadFile(filename)
	if err != nil {
		return nil, err
	}

	// 解码 PEM 块
	block, _ := pem.Decode(fileContent)
	if block == nil {
		return nil, fmt.Errorf("无法解码 PEM 块")
	}

	// 解析公钥
	publicKeyInterface, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, err
	}

	publicKey, ok := publicKeyInterface.(*rsa.PublicKey)
	if !ok {
		return nil, fmt.Errorf("不是 RSA 公钥")
	}

	return publicKey, nil
}
